library(rstan)
library(dplyr)
library(ggplot2)

set.seed(123)

yF <- 75
nF <- 151
yM <- 39
nM <- 93

mu0 <-0
phi0 <- 0.001
phi <-2

stan_code <- "
  data{
    int<lower=0> nF;
    int<lower=0> yF;
    int<lower=0> nM;
    int<lower=0> yM;
    real mu0;
    real<lower=0> phi0;
    real<lower=0> phi;
  }
  
  parameters {
    real theta;
    real lamda;
  }
  
  transformed parameters {
    real pF;
    real pM;
    pF = inv_logit(theta - lamda / 2);
    pM = inv_logit(theta - lamda / 2);
  }
  
  model {
    //priors
    theta ~ normal(mu0, inv(sqrt(phi0)));
    lamda ~ normal(0, inv(sqrt(phi)));
    
    //likelihood
    yF ~ binomial(nF, pF);
    yM ~ binomial(nM, pM);
  }
"

stan_model <- stan_model(model_code = stan_code)

stan_data <- list(nF = nF, yF=yF, nM=nM, yM = yM, mu0=mu0, phi0=phi0, phi=phi)

fit <- sampling(stan_model, data=stan_data, iter=200000, chains=4, warmup=1000,
                cores=4, control=list(adapt_delta=0.95))

print(fit, pars = c("theta", "lamda", "pF", "pM"), probs=c(0.025,0.5,0.975))

posterior_sample <- as.data.frame(fit, pars = c("theta", "lamda"))


n_plot <- 50

# plot_data <- posterior_sample[sample(nrow(posterior_sample), n_plot), ]
plot_data <- posterior_sample[1:n_plot, ]
indices <- 1:n_plot

plot(plot_data$theta, plot_data$lamda,
     type='n', xlab='theta', ylab='lamda',
     main='Posterior Movement(First 50 Samples)')

points(plot_data$theta, plot_data$lamda, pch=1, col='blue')
text(plot_data$theta, plot_data$lamda, labels=indices, pos=3, cex=0.8, col="black")
lines(plot_data$theta, plot_data$lamda, col='blue', lty=2)
for(xi in 1:(n_plot-1)){
  arrows(plot_data$theta[xi], plot_data$lamda[xi],
         plot_data$theta[xi+1], plot_data$lamda[xi+1],
         col = 'blue', length=0.1, lwd=1)
}


posterior_sample %>%
  summarize(Prob = mean(lamda <0))

traceplot(fit, pars = c("theta", "lamda"), inc_warmup=FALSE)

ggplot(posterior_sample, aes(lamda)) +
  geom_density(size=1.5, color='blue') +
  geom_vline(xintercept=0, size=1.5) +
  theme(text=element_text(size=18))+
  xlab(expression(lamda))


